import codecs
import re
import sys
import nltk
from collections import OrderedDict


event_path = sys.argv[1]
orderd_tgt_event = sys.argv[2]
stand_tgt_event = '/home/rickwwang/project_own/story_generation/data/test.wp_target_500.gensim.format.promptabs'
stand_tgt_story = '/home/rickwwang/project_own/story_generation/data/test.wp_target_500.gensim.format.story'

story_events = codecs.open(event_path, 'r', encoding='utf8').readlines()
tgt_event = codecs.open(orderd_tgt_event, 'r', encoding='utf8').readlines()
org_event = codecs.open(stand_tgt_event, 'r', encoding='utf8').readlines()
org_sotry = codecs.open(stand_tgt_story, 'r', encoding='utf8').readlines()
assert len(story_events) == len(tgt_event)
assert len(org_event) == len(org_sotry)
event_map = zip(org_event, range(len(org_event)))
event_map = dict(event_map)

result_tgt = event_path + '.tgt'
story_tgt = []
for i, story in enumerate(story_events):

    text = tgt_event[i]
    if i == len(story_events) - 1:
        text = text + '\n'
    org_index = event_map[text]
    story_tgt.append(org_sotry[org_index])

with codecs.open(result_tgt, 'w', encoding='utf8') as fout:
    fout.write(''.join(story_tgt))